#!/usr/bin/env python3

# Copyright © 2025 Wenze Wei
#
# This file is part of Pisces L1.
#
# Licensed under the Creative Commons Attribution-NonCommercial 4.0 International License (CC BY-NC 4.0).
# You may not use this file except in compliance with the License.
# Commercial use is strictly prohibited.
# You may obtain a copy of the License at
#
#     https://creativecommons.org/licenses/by-nc/4.0/
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import json
import pyarrow as pa
from datasets import Dataset, DatasetDict, load_dataset, load_from_disk

def arrow(args):
    """
    Convert between JSON files and Arrow format dataset.

    Args:
        args (argparse.Namespace): Command line arguments containing the following possible attributes:
            - json_dir (str): Directory path containing JSON files.
            - arrow_out (str): Output path to save the Arrow format dataset.
            - arrow_in (str): Input path of the Arrow format dataset.
            - json_out (str): Output path to save the JSON file.

    Returns:
        None: This function does not return any value except early returns in error cases.
    """
    # Check if the task is to convert JSON files to Arrow format
    if args.json_dir and args.arrow_out:
        # Get all JSON files in the specified directory
        json_files = [os.path.join(args.json_dir, f) for f in os.listdir(args.json_dir) if f.endswith('.json')]
        # Check if there are any JSON files in the directory
        if not json_files:
            print(f"❌\tNo .json files found in {args.json_dir}")
            return
        # Initialize an empty list to store all JSON data
        all_data = []
        # Iterate over each JSON file
        for jf in json_files:
            # Open the JSON file for reading
            with open(jf, 'r', encoding='utf-8') as f:
                # Read the file line by line
                for line in f:
                    try:
                        # Parse each line as JSON and add it to the data list
                        all_data.append(json.loads(line))
                    except Exception as e:
                        print(f"❌\tError parsing {jf}: {e}")
        # Check if any data was loaded from the JSON files
        if not all_data:
            print("❌\tNo data loaded from json files.")
            return
        # Create a dataset from the loaded data
        ds = Dataset.from_list(all_data)
        # Save the dataset to disk in Arrow format
        ds.save_to_disk(args.arrow_out)
        print(f"✅\tSaved {len(ds)} samples to {args.arrow_out}")
        return
    # Check if the task is to convert Arrow format dataset to JSON file
    elif args.arrow_in and args.json_out:
        # Check if the input Arrow file exists
        if not os.path.exists(args.arrow_in):
            print(f"❌\tArrow file not found: {args.arrow_in}")
            return
        # Load the dataset from disk
        ds = load_from_disk(args.arrow_in)
        # Open the output JSON file for writing
        with open(args.json_out, 'w', encoding='utf-8') as f:
            # Iterate over each item in the dataset
            for item in ds:
                # Write each item as a JSON line to the file
                f.write(json.dumps(item, ensure_ascii=False) + '\n')
        print(f"✅\tSaved {len(ds)} samples to {args.json_out}")
        return
    else:
        print("❌\tPlease specify either --json_dir + --arrow_out or --arrow_in + --json_out")
        print("For example：")
        print("\tpython manage.py arrow --json_dir ./jsons --arrow_out ./out.arrow")
        print("\tpython manage.py arrow --arrow_in ./in.arrow --json_out ./out.json")
